#!/usr/bin/env python3
"""
smartmontools MCP Server

面向 smartmontools 软件的 MCP Server，提供 smartd 服务管理和配置相关的命令行工具。

MCP tools 列表：
1. smartd_service_control - 控制 smartd 服务的启动、停止、重启、重载和状态查询
   参数：action (str): 操作类型，支持 start/stop/restart/reload/status/daemon-reload
   返回：统一的 JSON 结构，包含命令执行结果

2. smartd_config_modify - 修改 smartd 服务的配置
   参数：config_file (str): 配置文件路径
         old_pattern (str): 要替换的旧模式
         new_pattern (str): 替换的新模式
   返回：统一的 JSON 结构，包含命令执行结果
"""

from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import Dict, Any

mcp = FastMCP("Smartmontools MCP Server")

@mcp.tool()
def smartd_service_control(action: str) -> Dict[str, Any]:
    """
    控制 smartd 服务的各种操作
    
    参数:
        action: 操作类型，支持以下值：
                - start: 启动服务
                - stop: 停止服务  
                - restart: 重启服务
                - reload: 重载服务
                - status: 查看服务状态
                - daemon-reload: 重载 systemd 守护进程
    
    返回:
        JSON 结构:
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    valid_actions = ["start", "stop", "restart", "reload", "status", "daemon-reload"]
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} smartd.service",
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Invalid action: {action}. Valid actions are: {', '.join(valid_actions)}"
        }
    
    if action == "daemon-reload":
        command = ["systemctl", "daemon-reload"]
    else:
        command = ["systemctl", action, "smartd.service"]
    
    try:
        result = subprocess.run(command, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": " ".join(command),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(command),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command execution timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(command),
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def smartd_config_modify(config_file: str, old_pattern: str, new_pattern: str) -> Dict[str, Any]:
    """
    使用 sed 命令修改 smartd 服务的配置文件
    
    参数:
        config_file: 配置文件的完整路径
        old_pattern: 要替换的旧文本模式
        new_pattern: 替换的新文本模式
    
    返回:
        JSON 结构:
        {
            "success": bool,      # 操作是否成功
            "command": str,       # 执行的命令
            "exit_code": int,     # 退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误
        }
    """
    # 构建 sed 命令
    sed_command = ["sed", "-i", f"s\\{old_pattern}\\{new_pattern}\\g", config_file]
    
    try:
        result = subprocess.run(sed_command, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": " ".join(sed_command),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(sed_command),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command execution timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(sed_command),
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()